﻿/*
 * @author: wizardc
 */

using System;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;
using Random = UnityEngine.Random;

namespace Dou.Cluster
{
    /// <summary>
    /// 二分 k-means 算法
    /// </summary>
    public class BinarySearchKMeansCluster<T> where T : IKMeansClient
    {
        private T[] _centroids = new T[2];
        
        private List<List<T>> _clusterList = new List<List<T>>
        {
            new List<T>(),
            new List<T>()
        };
        
        public BinarySearchKMeansCluster(List<T> list)
        {
            // 质心寻找采用 k-means++ 算法思想
            if (list.Count < 2)
            {
                throw new Exception("聚类计算数量必须大于1");
            }
            // 随机取一个点为第一个质心
            var firstCentroid = list[Random.Range(0, list.Count)];
            _centroids[0] = firstCentroid;
            // 取距离第一个质心最远的一个点为第二个质心
            var maxDist = 0f;
            T maxDistObj = default;
            for (int i = 0; i < list.Count; i++)
            {
                var go = list[i];
                var dist = Vector3.Distance(firstCentroid.GetPosition(), go.GetPosition());
                if (dist > maxDist)
                {
                    maxDist = dist;
                    maxDistObj = go;
                }
            }
            _centroids[1] = maxDistObj;
        }

        public List<List<T>> clusterList => _clusterList;

        public void Cluster(List<T> list)
        {
            if (list.Count < 2)
            {
                throw new Exception("聚类计算数量必须大于1");
            }
            
            if (list.Count == 2)
            {
                _clusterList[0].Clear();
                _clusterList[1].Clear();
                
                _clusterList[0].Add(_centroids[0]);
                _clusterList[1].Add(_centroids[1]);
                return;
            }

            bool centroidsChanged;
            do
            {
                centroidsChanged = false;
                
                _clusterList[0].Clear();
                _clusterList[1].Clear();

                // 按距离分为2个聚落
                foreach (var obj in list)
                {
                    var closestClusterIndex = FindClosestCluster(obj);
                    _clusterList[closestClusterIndex].Add(obj);
                }

                // 重算质心
                for (var i = 0; i < 2; i++)
                {
                    if (_clusterList[i].Count == 0)
                    {
                        break;
                    }

                    var newCentroid = CalculateCentroid(_clusterList[i]);
                    if (!SamePoint(_centroids[i], newCentroid))
                    {
                        _centroids[i] = newCentroid;
                        centroidsChanged = true;
                    }
                }
            } while (centroidsChanged);
        }
        
        private int FindClosestCluster(T obj)
        {
            var distanceToFirst = Vector3.Distance(obj.GetPosition(), _centroids[0].GetPosition());
            var distanceToSecond = Vector3.Distance(obj.GetPosition(), _centroids[1].GetPosition());
            return distanceToFirst < distanceToSecond ? 0 : 1;
        }
        
        private T CalculateCentroid(List<T> cluster)
        {
            var averagePosition = Vector3.zero;
            foreach (var obj in cluster)
            {
                averagePosition += obj.GetPosition();
            }
            averagePosition /= cluster.Count;
            // 找到距离质心最近的对象返回
            return cluster.OrderBy(obj => Vector3.Distance(obj.GetPosition(), averagePosition)).First();
        }

        private bool SamePoint(T lhs, T rhs)
        {
            return lhs.GetPosition() == rhs.GetPosition();
        }
    }
}
